import time
import tkinter as tk
from PIL import Image, ImageDraw, ImageOps
import numpy as np
from main import load_saved_model
import torch

# 加载预训练的MNIST模型（需要提前准备模型文件）
model = load_saved_model()  # 请确保模型文件存在


class DrawingApp:
    def __init__(self, root):
        self.root = root
        self.root.title("手写数字识别")

        # 设置画布尺寸
        self.canvas_width = 280
        self.canvas_height = 280
        self.scale_factor = 10  # 从MNIST的28x28放大到280x280

        # 创建画布
        self.canvas = tk.Canvas(root, width=self.canvas_width, height=self.canvas_height, bg='black')
        self.canvas.pack()

        # 设置按钮框架
        button_frame = tk.Frame(root)
        button_frame.pack(fill=tk.X, padx=5, pady=5)

        # 创建按钮
        self.clear_button = tk.Button(button_frame, text="清除", command=self.clear_canvas)
        self.clear_button.pack(side=tk.LEFT, padx=5)

        self.recognize_button = tk.Button(button_frame, text="识别", command=self.recognize_digit)
        self.recognize_button.pack(side=tk.LEFT, padx=5)

        self.save_button = tk.Button(button_frame, text="保存", command=self.save_image)
        self.save_button.pack(side=tk.LEFT, padx=5)

        # 创建结果显示标签
        self.result_label = tk.Label(root, text="", font=('Arial', 12))
        self.result_label.pack()

        # 初始化绘图变量
        self.last_x = None
        self.last_y = None
        self.drawing_data = []

        # 绑定事件
        self.canvas.bind("<Button-1>", self.start_drawing)
        self.canvas.bind("<B1-Motion>", self.draw)

    def start_drawing(self, event):
        self.last_x = event.x
        self.last_y = event.y

    def draw(self, event):
        x, y = event.x, event.y
        if self.last_x and self.last_y:
            # 在画布上绘制
            self.canvas.create_line(self.last_x, self.last_y, x, y,
                                    width=20, fill='white', capstyle=tk.ROUND, smooth=True)
            # 记录绘制数据
            self.drawing_data.append(('line', (self.last_x, self.last_y, x, y), 'white', 20))
        self.last_x = x
        self.last_y = y

    def clear_canvas(self):
        self.canvas.delete("all")
        self.drawing_data = []
        self.result_label.config(text="")

    def save_image(self):
        # 创建空白图像
        img = Image.new("L", (self.canvas_width, self.canvas_height), 0)
        draw = ImageDraw.Draw(img)

        # 重绘所有记录的操作
        for item in self.drawing_data:
            if item[0] == 'line':
                x1, y1, x2, y2 = item[1]
                draw.line([x1, y1, x2, y2], fill=255, width=item[3])

        # 保存原始图像
        img.save("handwritten_digit.png")
        self.result_label.config(text="Saved!")

    def recognize_digit(self):
        # 创建临时图像进行识别
        img = Image.new("L", (self.canvas_width, self.canvas_height), 0)
        draw = ImageDraw.Draw(img)

        for item in self.drawing_data:
            if item[0] == 'line':
                x1, y1, x2, y2 = item[1]
                draw.line([x1, y1, x2, y2], fill=255, width=item[3])
        # display
        import matplotlib.pyplot as plt
        from main import _transform
        img = _transform(img)
        plt.imshow(img.squeeze(), cmap='gray')
        plt.show()
        img = img.expand(1, -1, -1, -1)
        outputs = model(img)
        probs = torch.exp(outputs)
        output = outputs[0]

        confidence = probs.squeeze()

        predicted_class = torch.argmax(output)
        confidence = confidence[predicted_class]
        print(f"Image result: {predicted_class}, Confidence: {confidence}")
        print(img.shape)
        self.result_label.config(text=f"识别结果：{predicted_class}（置信度：{confidence:.2f}）")


if __name__ == "__main__":
    root = tk.Tk()
    app = DrawingApp(root)
    root.mainloop()
